package com.feiwo.nonego.rpc.loadbalancer.consistent.hash;

import com.feiwo.nonego.rpc.loadbalancer.api.ServiceLoadBalancer;
import com.feiwo.nonego.rpc.protocol.meta.ServiceMeta;
import com.feiwo.nonego.rpc.spi.annotation.SPIClass;
import org.apache.curator.x.discovery.ServiceInstance;

import java.util.List;
import java.util.Objects;
import java.util.TreeMap;

/**
 *
 *
 * @author feiwo 2025-08-19
 */
@SPIClass
public class ZKConsistentHashLoadBalancer implements ServiceLoadBalancer<ServiceInstance<ServiceMeta>> {

    private static final int VIRTUAL_NODE_NUM = 10;
    private static final String VIRTUAL_NODE_SPLIT = "#";

    @Override
    public ServiceInstance<ServiceMeta> select(List<ServiceInstance<ServiceMeta>> servers, int hashCode, String sourceIp) {
        var ring = makeConsistentHashRing(servers);
        return allocateNode(ring, hashCode);
    }

    private ServiceInstance<ServiceMeta> allocateNode(TreeMap<Integer, ServiceInstance<ServiceMeta>> ring, int hashCode) {
        var entry = ring.ceilingEntry(hashCode);
        if (Objects.isNull(entry)) {
            entry = ring.firstEntry();
        }
        if (Objects.isNull(entry)) {
            throw new IllegalStateException("not discover useful service, please register service in registry center.");
        }
        return entry.getValue();
    }

    private TreeMap<Integer, ServiceInstance<ServiceMeta>> makeConsistentHashRing(List<ServiceInstance<ServiceMeta>> services) {
        var ring = new TreeMap<Integer, ServiceInstance<ServiceMeta>>();
        for (ServiceInstance<ServiceMeta> instance : services) {
            for (int i = 0; i < VIRTUAL_NODE_NUM; i++) {
                ring.put((buildKey(instance)+ VIRTUAL_NODE_SPLIT+ i).hashCode(), instance);
            }
        }
        return ring;
    }

    private static String buildKey(ServiceInstance<ServiceMeta> instance) {
        var payload = instance.getPayload();
        return String.join(":", payload.getServiceName(), payload.getServiceVersion(), payload.getServiceGroup());
    }
}
